package com.websocket.utils;

import lombok.extern.slf4j.Slf4j;

import javax.websocket.Session;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @author qinwei
 * @date 2024/11/02 周六 16:47
 * @description
 */
@Slf4j
public class WebsocketUtil {

    private static final Map<String, Session> ONLINE_SESSION = new ConcurrentHashMap<>();

    /**
     * 添加session
     */
    public static void addSession(String userId, Session session){
        // 一个用户只允许一个session链接
        ONLINE_SESSION.putIfAbsent(userId, session);
        log.info("User [{}] connected. Total online users: {}", userId, ONLINE_SESSION.size());
    }

    /**
     * 移除session
     */
    public static void removeSession(String userId){
        ONLINE_SESSION.remove(userId);
        log.info("User [{}] disconnected. Total online users: {}", userId, ONLINE_SESSION.size());
    }

    /**
     * 给单个用户推送消息
     */
    public static void sendMessage(String userId, String message){
        Session session = ONLINE_SESSION.get(userId);
        if(session == null){
            log.warn("Session for user [{}] not found", userId);
            return;
        }
        sendMessage(session, message);
    }

    /**
     * 发送文字
     * @param session
     * @param message
     */
    public static void sendMessage(Session session, String message) {
        if (session != null) {
            session.getAsyncRemote().sendText(message);
        }
    }

    /**
     * 发送文件流
     * @param imageData
     * @param session
     * @throws IOException
     */
    public void sendImage(byte[] imageData, Session session) throws IOException {
        // 发送图片数据
        session.getBasicRemote().sendBinary(ByteBuffer.wrap(imageData));
    }

    /**
     * 给所有用户发消息
     */
    public static void sendMessageForAll(String message) {
        ONLINE_SESSION.forEach((userId, session) -> {
            CompletableFuture.runAsync(() -> sendMessage(session, message))
                    .exceptionally(ex -> {
                        log.error("Failed to send message to user [{}]: {}", userId, ex.getMessage());
                        return null;
                    });
        });
    }

    /**
     * 给指定的多个用户推送消息
     */
    public static void sendMessageForUsers(Set<String> userIds, String message) {
        userIds.forEach(userId -> {
            Session session = ONLINE_SESSION.get(userId);
            if (session == null) {
                log.warn("Session for user [{}] not found", userId);
                return;
            }
            CompletableFuture.runAsync(() -> sendMessage(session, message))
                    .exceptionally(ex -> {
                        log.error("Failed to send message to user [{}]: {}", userId, ex.getMessage());
                        return null;
                    });
        });
    }
}
